import json
import uuid
import threading
import time
import paramiko
from channels.generic.websocket import AsyncWebsocketConsumer
from channels.db import database_sync_to_async
from django.contrib.auth import get_user_model
from django.utils import timezone
import asyncio
import logging

logger = logging.getLogger(__name__)


class SSHConsumer(AsyncWebsocketConsumer):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.ssh_client = None
        self.ssh_channel = None
        self.host = None
        self.user = None
        self.connection_log = None
        self.session_id = None
        self.is_connected = False
        self.last_activity = time.time()
        
    async def connect(self):
        """WebSocket连接建立"""
        self.host_id = self.scope['url_route']['kwargs']['host_id']
        self.user = self.scope['user']
        self.session_id = str(uuid.uuid4())
        
        if self.user.is_anonymous:
            await self.close()
            return
            
        # 获取主机信息
        self.host = await self.get_host(self.host_id)
        if not self.host:
            await self.close()
            return
            
        await self.accept()
        
        # 创建连接日志
        self.connection_log = await self.create_connection_log()
        
        # 发送连接状态
        await self.send(text_data=json.dumps({
            'type': 'status',
            'message': f'正在连接到 {self.host.hostname} ({self.host.ip_address})...',
            'status': 'connecting'
        }))
        
        # 异步建立SSH连接
        asyncio.create_task(self.establish_ssh_connection())
    
    async def disconnect(self, close_code):
        """WebSocket连接断开"""
        await self.cleanup_ssh_connection()
        if self.connection_log:
            await self.update_connection_log_status('disconnected')

        # 检查并更新主机状态
        await self.check_and_update_host_status_on_disconnect()
    
    async def receive(self, text_data):
        """接收WebSocket消息"""
        try:
            data = json.loads(text_data)
            message_type = data.get('type')
            
            if message_type == 'command':
                await self.handle_command(data.get('data', ''))
            elif message_type == 'resize':
                await self.handle_resize(data.get('cols', 80), data.get('rows', 24))
            elif message_type == 'ping':
                # 检查SSH连接状态并处理心跳
                await self.handle_heartbeat_ping(data)
                
        except json.JSONDecodeError:
            logger.error(f"Invalid JSON received: {text_data}")
        except Exception as e:
            logger.error(f"Error handling message: {e}")
    
    async def establish_ssh_connection(self):
        """建立SSH连接"""
        try:
            # 更新主机状态为检测中
            await self.update_host_status('checking')

            # 在线程池中执行SSH连接
            loop = asyncio.get_event_loop()
            await loop.run_in_executor(None, self._create_ssh_connection)

            if self.is_connected:
                await self.send(text_data=json.dumps({
                    'type': 'status',
                    'message': f'已连接到 {self.host.hostname}',
                    'status': 'connected'
                }))
                await self.update_connection_log_status('connected')

                # 更新主机状态为在线
                await self.update_host_status('online')

                # 启动数据读取任务
                asyncio.create_task(self.read_ssh_output())
            else:
                await self.send(text_data=json.dumps({
                    'type': 'status',
                    'message': 'SSH连接失败',
                    'status': 'failed'
                }))
                await self.update_connection_log_status('failed')

                # 更新主机状态为异常
                await self.update_host_status('error')

        except Exception as e:
            logger.error(f"SSH connection error: {e}")
            await self.send(text_data=json.dumps({
                'type': 'status',
                'message': f'连接失败: {str(e)}',
                'status': 'failed'
            }))
            await self.update_connection_log_status('failed', str(e))

            # 更新主机状态为异常
            await self.update_host_status('error')
    
    def _create_ssh_connection(self):
        """创建SSH连接（同步方法）"""
        try:
            self.ssh_client = paramiko.SSHClient()
            self.ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            
            # 连接到主机
            self.ssh_client.connect(
                hostname=self.host.ip_address,
                port=self.host.ssh_port or 22,
                username=self.host.ssh_username or 'root',
                password=self.host.ssh_password,
                timeout=10
            )
            
            # 创建交互式shell
            self.ssh_channel = self.ssh_client.invoke_shell(
                term='xterm-256color',
                width=80,
                height=24
            )
            
            self.is_connected = True
            
        except Exception as e:
            logger.error(f"SSH connection failed: {e}")
            self.is_connected = False
            raise e
    
    async def read_ssh_output(self):
        """读取SSH输出"""
        loop = asyncio.get_event_loop()
        
        while self.is_connected and self.ssh_channel:
            try:
                # 在线程池中读取数据
                data = await loop.run_in_executor(None, self._read_ssh_data)
                
                if data:
                    await self.send(text_data=json.dumps({
                        'type': 'output',
                        'data': data
                    }))
                    self.last_activity = time.time()
                else:
                    # 没有数据，短暂休眠
                    await asyncio.sleep(0.1)
                    
            except Exception as e:
                logger.error(f"Error reading SSH output: {e}")
                break
    
    def _read_ssh_data(self):
        """读取SSH数据（同步方法）"""
        try:
            if self.ssh_channel and self.ssh_channel.recv_ready():
                data = self.ssh_channel.recv(1024)
                return data.decode('utf-8', errors='ignore')
        except Exception as e:
            logger.error(f"Error reading SSH data: {e}")
        return None
    
    async def handle_command(self, command):
        """处理命令输入"""
        if not self.is_connected or not self.ssh_channel:
            return
            
        try:
            loop = asyncio.get_event_loop()
            await loop.run_in_executor(None, self._send_command, command)
            self.last_activity = time.time()
            
            # 记录命令日志
            if command.strip():
                await self.log_command(command.strip())
                
        except Exception as e:
            logger.error(f"Error sending command: {e}")
    
    def _send_command(self, command):
        """发送命令（同步方法）"""
        try:
            if self.ssh_channel:
                self.ssh_channel.send(command.encode('utf-8'))
        except Exception as e:
            logger.error(f"Error sending command: {e}")
    
    async def handle_resize(self, cols, rows):
        """处理终端大小调整"""
        if self.ssh_channel:
            try:
                loop = asyncio.get_event_loop()
                await loop.run_in_executor(None, self._resize_terminal, cols, rows)
                
                # 更新连接日志中的终端大小
                if self.connection_log:
                    await self.update_terminal_size(f"{cols}x{rows}")
                    
            except Exception as e:
                logger.error(f"Error resizing terminal: {e}")
    
    def _resize_terminal(self, cols, rows):
        """调整终端大小（同步方法）"""
        try:
            if self.ssh_channel:
                self.ssh_channel.resize_pty(width=cols, height=rows)
        except Exception as e:
            logger.error(f"Error resizing terminal: {e}")
    
    async def cleanup_ssh_connection(self):
        """清理SSH连接"""
        self.is_connected = False
        
        if self.ssh_channel:
            try:
                self.ssh_channel.close()
            except:
                pass
            self.ssh_channel = None
            
        if self.ssh_client:
            try:
                self.ssh_client.close()
            except:
                pass
            self.ssh_client = None
    
    @database_sync_to_async
    def get_host(self, host_id):
        """获取主机信息"""
        try:
            from .models import Host
            return Host.objects.get(id=host_id)
        except Host.DoesNotExist:
            return None
    
    @database_sync_to_async
    def create_connection_log(self):
        """创建连接日志"""
        from .models import SSHConnectionLog
        return SSHConnectionLog.objects.create(
            host=self.host,
            user=self.user,
            source_ip=self.scope.get('client', [''])[0],
            user_agent=dict(self.scope.get('headers', {})).get(b'user-agent', b'').decode(),
            session_id=self.session_id,
            status='connecting'
        )
    
    @database_sync_to_async
    def update_connection_log_status(self, status, error_message=''):
        """更新连接日志状态"""
        if self.connection_log:
            self.connection_log.status = status
            if error_message:
                self.connection_log.error_message = error_message
            if status == 'disconnected':
                self.connection_log.disconnect_time = timezone.now()
            self.connection_log.save()
    
    @database_sync_to_async
    def update_terminal_size(self, size):
        """更新终端大小"""
        if self.connection_log:
            self.connection_log.terminal_size = size
            self.connection_log.save()
    
    @database_sync_to_async
    def log_command(self, command):
        """记录命令日志"""
        if self.connection_log:
            from .models import SSHCommandLog
            SSHCommandLog.objects.create(
                connection=self.connection_log,
                command=command
            )

    async def handle_heartbeat_ping(self, data):
        """处理心跳ping消息"""
        try:
            # 检查SSH连接是否仍然活跃
            if self.is_connected and self.ssh_client and self.ssh_channel:
                # 尝试发送一个简单的命令来测试连接
                try:
                    # 在线程池中执行连接测试
                    loop = asyncio.get_event_loop()
                    connection_ok = await loop.run_in_executor(None, self._test_ssh_connection)

                    if connection_ok:
                        # 连接正常，发送pong响应
                        await self.send(text_data=json.dumps({
                            'type': 'pong',
                            'timestamp': data.get('timestamp', time.time()),
                            'status': 'connected'
                        }))

                        # 更新主机状态为在线
                        await self.update_host_status('online')
                    else:
                        # 连接异常，发送错误响应
                        await self.send(text_data=json.dumps({
                            'type': 'pong',
                            'timestamp': data.get('timestamp', time.time()),
                            'status': 'disconnected',
                            'error': 'SSH connection lost'
                        }))

                        # 更新主机状态为离线
                        await self.update_host_status('offline')

                        # 标记连接为断开
                        self.is_connected = False

                except Exception as e:
                    logger.error(f"Heartbeat connection test failed: {e}")

                    # 发送错误响应
                    await self.send(text_data=json.dumps({
                        'type': 'pong',
                        'timestamp': data.get('timestamp', time.time()),
                        'status': 'error',
                        'error': str(e)
                    }))

                    # 更新主机状态为离线
                    await self.update_host_status('offline')
                    self.is_connected = False
            else:
                # 连接已断开，发送断开响应
                await self.send(text_data=json.dumps({
                    'type': 'pong',
                    'timestamp': data.get('timestamp', time.time()),
                    'status': 'disconnected',
                    'error': 'No active SSH connection'
                }))

                # 更新主机状态为离线
                await self.update_host_status('offline')

        except Exception as e:
            logger.error(f"Error handling heartbeat ping: {e}")

            # 发送错误响应
            await self.send(text_data=json.dumps({
                'type': 'pong',
                'timestamp': data.get('timestamp', time.time()),
                'status': 'error',
                'error': str(e)
            }))

    def _test_ssh_connection(self):
        """测试SSH连接是否仍然活跃（同步方法）"""
        try:
            if not self.ssh_client or not self.ssh_channel:
                return False

            # 发送一个简单的命令来测试连接
            self.ssh_channel.send('echo heartbeat_test\n')

            # 等待响应（短时间）
            import select
            ready, _, _ = select.select([self.ssh_channel], [], [], 2.0)  # 2秒超时

            if ready:
                # 有数据可读，连接正常
                data = self.ssh_channel.recv(1024)
                return len(data) > 0
            else:
                # 超时，连接可能有问题
                return False

        except Exception as e:
            logger.error(f"SSH connection test error: {e}")
            return False

    async def update_host_status(self, status):
        """更新主机状态的通用方法"""
        if hasattr(self, 'host') and self.host:
            try:
                from django.utils import timezone

                # 异步更新主机状态
                @database_sync_to_async
                def update_status():
                    # 只有在状态真正改变时才更新
                    if self.host.status != status:
                        old_status = self.host.status
                        self.host.status = status
                        self.host.last_check_time = timezone.now()
                        self.host.save(update_fields=['status', 'last_check_time'])
                        logger.info(f"Host {self.host.hostname} status changed from {old_status} to {status}")
                        return True
                    else:
                        # 即使状态没变，也更新检查时间
                        self.host.last_check_time = timezone.now()
                        self.host.save(update_fields=['last_check_time'])
                        return False

                status_changed = await update_status()

                # 如果状态发生变化，可以在这里添加通知逻辑
                if status_changed:
                    await self.broadcast_host_status_change(status)

            except Exception as e:
                logger.error(f"Error updating host status to {status}: {e}")

    async def broadcast_host_status_change(self, new_status):
        """广播主机状态变化（可用于实时更新前端）"""
        try:
            from .host_status_consumer import broadcast_host_status_update

            # 广播主机状态更新
            await broadcast_host_status_update(
                host_id=self.host.id,
                hostname=self.host.hostname,
                status=new_status,
                last_check_time=self.host.last_check_time
            )

        except Exception as e:
            logger.error(f"Error broadcasting host status change: {e}")

    async def check_and_update_host_status_on_disconnect(self):
        """检查并更新主机断开连接时的状态"""
        if hasattr(self, 'host') and self.host:
            try:
                from .models import SSHConnectionLog
                from django.utils import timezone

                # 检查是否还有其他活跃的SSH连接
                @database_sync_to_async
                def check_active_connections():
                    active_connections = SSHConnectionLog.objects.filter(
                        host=self.host,
                        status='connected'
                    ).count()

                    # 如果没有其他活跃连接，将主机状态设为离线
                    if active_connections == 0:
                        self.host.status = 'offline'
                        self.host.last_check_time = timezone.now()
                        self.host.save(update_fields=['status', 'last_check_time'])
                        return True
                    return False

                updated = await check_active_connections()
                if updated:
                    logger.debug(f"Updated host {self.host.hostname} status to offline (no active connections)")

            except Exception as e:
                logger.error(f"Error checking host status on disconnect: {e}")
